import asyncio
import json
from typing import List, Dict, Any
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.future import select
from sqlalchemy import and_
from loguru import logger
from datetime import datetime
from fastapi.concurrency import run_in_threadpool

from app.models.db import AnalysisDimension, DataSet, async_session_maker
from app.ai.ai_getdata import AIGetData
from app.services.dataset_service import DatasetService
from app.datatk.data_help import DataHelp
from config import AI_API_KEY, AI_BASEURL


class TaskExecutor:
    """任务执行器类，负责按顺序执行分析维度中的任务"""
    
    # 任务状态常量
    STATUS_PENDING = "pending"      # 未开始
    STATUS_IN_PROGRESS = "in_progress"  # 进行中
    STATUS_COMPLETED = "completed"  # 已完成
    STATUS_FAILED = "failed"        # 失败
    
    def __init__(self):
        self.ai_advisor = AIGetData(AI_BASEURL, AI_API_KEY)
    
    async def get_tasks_status_by_message(
        self, 
        message_id: str, 
        session: AsyncSession
    ) -> Dict[str, Any]:
        """
        根据消息ID获取任务执行状况
        
        参数:
            message_id (str): 消息ID
            session (AsyncSession): 数据库会话
            
        返回:
            Dict[str, Any]: 任务执行状况
        """
        try:
            # 查询该消息关联的所有任务
            tasks_query = select(AnalysisDimension).where(
                AnalysisDimension.report_id == message_id,
                AnalysisDimension.dimension_type == "task"
            ).order_by(AnalysisDimension.dimension_order, AnalysisDimension.created_at)
            
            tasks_result = await session.execute(tasks_query)
            tasks = tasks_result.scalars().all()
            
            if not tasks:
                return {
                    "status": "no_tasks",
                    "message": "未找到相关任务",
                    "total_tasks": 0,
                    "tasks": [],
                    "all_query_results": []
                }
            
            # 统计各状态任务数量
            status_counts = {
                self.STATUS_PENDING: 0,
                self.STATUS_IN_PROGRESS: 0,
                self.STATUS_COMPLETED: 0,
                self.STATUS_FAILED: 0
            }
            
            # 构建任务详情列表
            task_details = []
            # 收集数据查询任务的结果（参考task_breakdown接口的逻辑）
            all_query_results = []
            
            for task in tasks:
                task_status = task.status or self.STATUS_PENDING
                status_counts[task_status] = status_counts.get(task_status, 0) + 1
                
                task_detail = {
                    "task_id": task.id,
                    "task_name": task.dimension_name,
                    "task_type": task.analysis_method,
                    "task_order": task.dimension_order or 0,
                    "status": task_status,
                    "created_at": task.created_at.isoformat() if task.created_at else None,
                    "has_sql": bool(task.analysis_sql),
                    "has_result": bool(task.analysis_result),
                    "has_data": bool(task.json_data)
                }
                task_details.append(task_detail)
                
                # 收集已完成的数据查询任务的结果数据
                if (task.analysis_method == "数据查询" and 
                    task_status == self.STATUS_COMPLETED and 
                    task.json_data):
                    try:
                        # 解析任务要求以获取数据说明
                        task_requirements = {}
                        datadec = ""
                        
                        if task.dimension_description:
                            try:
                                task_requirements = json.loads(task.dimension_description)
                                datadec = task_requirements.get("说明", "")
                            except json.JSONDecodeError:
                                datadec = task.dimension_description  # 如果不是JSON，直接使用原文本
                        
                        # 解析json_data
                        json_data = []
                        try:
                            json_data = json.loads(task.json_data)
                        except json.JSONDecodeError:
                            json_data = []
                        
                        # 构建结果对象（与task_breakdown接口保持一致的格式）
                        result_item = {
                            "data": json_data,
                            "datadec": datadec,
                            "task_name": task.dimension_name,
                            "task_id": task.id
                        }
                        
                        all_query_results.append(result_item)
                        
                    except Exception as e:
                        logger.error(f"处理任务 {task.id} 的结果时出错: {str(e)}")
                        continue
            
            # 判断整体执行状态
            total_tasks = len(tasks)
            if status_counts[self.STATUS_FAILED] > 0:
                overall_status = "has_failures"
            elif status_counts[self.STATUS_COMPLETED] == total_tasks:
                overall_status = "all_completed"
            elif status_counts[self.STATUS_IN_PROGRESS] > 0:
                overall_status = "in_progress"
            else:
                overall_status = "pending"
            
            return {
                "status": overall_status,
                "message": f"共 {total_tasks} 个任务",
                "total_tasks": total_tasks,
                "status_counts": status_counts,
                "tasks": task_details,
                "all_query_results": all_query_results,  # 添加数据查询结果
                "progress_percentage": round((status_counts[self.STATUS_COMPLETED] / total_tasks) * 100, 1) if total_tasks > 0 else 0
            }
            
        except Exception as e:
            return {
                "status": "error",
                "message": f"获取任务状况失败: {str(e)}",
                "total_tasks": 0,
                "tasks": [],
                "all_query_results": []
            }
    
    async def execute_tasks_by_message(
        self, 
        message_id: str, 
        session: AsyncSession
    ) -> Dict[str, Any]:
        """
        根据消息ID执行相关的任务集
        
        参数:
            message_id (str): 消息ID
            session (AsyncSession): 数据库会话
            
        返回:
            Dict[str, Any]: 执行结果
        """
        try:
            # 1. 查询该消息关联的所有任务
            tasks_query = select(AnalysisDimension).where(
                AnalysisDimension.report_id == message_id,
                AnalysisDimension.dimension_type == "task",
                AnalysisDimension.status.in_(["pending"])
            ).order_by(AnalysisDimension.dimension_order, AnalysisDimension.created_at)
            
            tasks_result = await session.execute(tasks_query)
            tasks = tasks_result.scalars().all()
            
            if not tasks:
                return {
                    "status": False,
                    "message": "未找到相关任务",
                }
            
            
            # 2. 按任务组号分组
            task_groups = self._group_tasks_by_order(tasks)
            
            # 3. 按组号顺序执行任务
            executed_count = 0
            failed_count = 0
            
            for group_order in sorted(task_groups.keys()):
                group_tasks = task_groups[group_order]
                
                # 同一组内的任务可以并行执行
                group_results = await self._execute_task_group(group_tasks)
                
                # 统计执行结果
                for result in group_results:
                    if result.get("success", False):
                        executed_count += 1
                    else:
                        failed_count += 1
            
            return {
                "status": True,
                "message": f"任务执行完成，成功 {executed_count} 个，失败 {failed_count} 个"
            }
            
        except Exception as e:
            return {
                "status": False,
                "message": f"执行任务失败: {str(e)}"
            }
    
    def _group_tasks_by_order(self, tasks: List[AnalysisDimension]) -> Dict[int, List[AnalysisDimension]]:
        """
        按任务组号对任务进行分组
        
        参数:
            tasks (List[AnalysisDimension]): 任务列表
            
        返回:
            Dict[int, List[AnalysisDimension]]: 按组号分组的任务字典
        """
        task_groups = {}
        
        for task in tasks:
            group_order = task.dimension_order or 0
            if group_order not in task_groups:
                task_groups[group_order] = []
            task_groups[group_order].append(task)
        
        return task_groups
    
    async def _execute_task_group(
        self, 
        tasks: List[AnalysisDimension]
    ) -> List[Dict[str, Any]]:
        """
        执行一组任务（并行执行）
        
        参数:
            tasks (List[AnalysisDimension]): 任务列表
            
        返回:
            List[Dict[str, Any]]: 执行结果列表
        """
        # 创建任务协程列表，每个任务使用独立的数据库会话
        task_coroutines = []
        
        for task in tasks:
            coroutine = self._execute_single_task_with_session(task)
            task_coroutines.append(coroutine)
        
        # 并行执行所有任务
        results = await asyncio.gather(*task_coroutines, return_exceptions=True)
        
        # 处理结果，将异常转换为错误结果
        processed_results = []
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                processed_results.append({
                    "success": False,
                    "task_id": tasks[i].id,
                    "error": str(result)
                })
            else:
                processed_results.append(result)
        
        return processed_results
    
    async def _execute_single_task_with_session(
        self, 
        task: AnalysisDimension
    ) -> Dict[str, Any]:
        """
        使用独立数据库会话执行单个任务
        
        参数:
            task (AnalysisDimension): 任务对象
            
        返回:
            Dict[str, Any]: 执行结果
        """
        task_session = None
        try:
            # 为每个任务创建独立的数据库会话
            task_session = async_session_maker()
            
            try:
                # 重新查询任务对象，确保在当前会话中
                task_query = select(AnalysisDimension).where(AnalysisDimension.id == task.id)
                task_result = await task_session.execute(task_query)
                current_task = task_result.scalar_one_or_none()
                
                if not current_task:
                    return {
                        "success": False,
                        "task_id": task.id,
                        "error": "任务不存在"
                    }
                
                # 执行任务
                result = await self._execute_single_task(current_task, task_session)
                
                # 确保所有事务都已提交
                if task_session.in_transaction():
                    await task_session.commit()
                
                return result
                
            except Exception as e:
                # 发生异常时回滚事务
                if task_session.in_transaction():
                    await task_session.rollback()
                
                return {
                    "success": False,
                    "task_id": task.id,
                    "error": str(e)
                }
                
        except Exception as e:
            return {
                "success": False,
                "task_id": task.id,
                "error": f"会话创建失败: {str(e)}"
            }
        finally:
            # 确保会话被正确关闭
            if task_session is not None:
                try:
                    await task_session.close()
                except Exception as close_error:
                    logger.error(f"关闭任务 {task.id} 的数据库会话时发生错误: {str(close_error)}")
    
    async def _execute_single_task(
        self, 
        task: AnalysisDimension, 
        session: AsyncSession
    ) -> Dict[str, Any]:
        """
        执行单个任务
        
        参数:
            task (AnalysisDimension): 任务对象
            session (AsyncSession): 数据库会话
            
        返回:
            Dict[str, Any]: 执行结果
        """
        try:            
            # 更新任务状态为进行中
            task.status = self.STATUS_IN_PROGRESS
            task.updated_at = datetime.now()
            session.add(task)
            await session.commit()
            
            # 第一步：判断是否为"未知任务"
            if task.analysis_method == "未知任务":
                task.status = self.STATUS_COMPLETED
                task.updated_at = datetime.now()
                session.add(task)
                await session.commit()
                
                return {
                    "success": True,
                    "task_id": task.id,
                    "message": "跳过未知任务",
                    "skipped": True
                }
            
            # 第二步：根据任务类型执行相应操作
            if task.analysis_method == "数据查询":
                return await self._execute_data_query_task(task, session)
            elif task.analysis_method == "图表制作":
                task.status = self.STATUS_COMPLETED
                task.updated_at = datetime.now()
                session.add(task)
                await session.commit()
                
                return {
                    "success": True,
                    "task_id": task.id,
                    "message": "图表制作任务暂不处理",
                    "skipped": True
                }
            elif task.analysis_method == "报告生成":
                task.status = self.STATUS_COMPLETED
                task.updated_at = datetime.now()
                session.add(task)
                await session.commit()
                
                return {
                    "success": True,
                    "task_id": task.id,
                    "message": "报告生成任务暂不处理",
                    "skipped": True
                }
            else:
                task.status = self.STATUS_FAILED
                task.updated_at = datetime.now()
                session.add(task)
                await session.commit()
                
                return {
                    "success": False,
                    "task_id": task.id,
                    "error": f"未知任务类型: {task.analysis_method}"
                }
                
        except Exception as e:
            # 更新任务状态为失败
            try:
                if session.in_transaction():
                    await session.rollback()
                
                task.status = self.STATUS_FAILED
                task.updated_at = datetime.now()
                session.add(task)
                await session.commit()
            except Exception as commit_error:
                try:
                    if session.in_transaction():
                        await session.rollback()
                except:
                    pass
            
            return {
                "success": False,
                "task_id": task.id,
                "error": str(e)
            }
    
    async def _execute_data_query_task(
        self, 
        task: AnalysisDimension, 
        session: AsyncSession
    ) -> Dict[str, Any]:
        """
        执行数据查询任务
        
        参数:
            task (AnalysisDimension): 任务对象
            session (AsyncSession): 数据库会话
            
        返回:
            Dict[str, Any]: 执行结果
        """
        try:
            # 解析任务要求
            task_requirements = {}
            if task.dimension_description:
                try:
                    task_requirements = json.loads(task.dimension_description)
                except json.JSONDecodeError:
                    await self._update_task_status_safely(task, self.STATUS_FAILED, session)
                    return {
                        "success": False,
                        "task_id": task.id,
                        "error": "任务要求格式不正确"
                    }
            
            # 获取数据集信息
            if not task.dataset_id or task.dataset_id == "0":
                await self._update_task_status_safely(task, self.STATUS_FAILED, session)
                return {
                    "success": False,
                    "task_id": task.id,
                    "error": "任务未关联数据集"
                }
            
            # 查询数据集
            dataset_query = select(DataSet).where(DataSet.id == task.dataset_id)
            dataset_result = await session.execute(dataset_query)
            dataset = dataset_result.scalar_one_or_none()
            
            if not dataset:
                await self._update_task_status_safely(task, self.STATUS_FAILED, session)
                return {
                    "success": False,
                    "task_id": task.id,
                    "error": f"数据集不存在: {task.dataset_id}"
                }
            
            # 获取数据集结构信息
            dataset_schema = {}
            if dataset.dataset_summary:
                try:
                    dataset_schema = json.loads(dataset.dataset_summary)
                except json.JSONDecodeError:
                    await self._update_task_status_safely(task, self.STATUS_FAILED, session)
                    return {
                        "success": False,
                        "task_id": task.id,
                        "error": "数据集结构信息格式不正确"
                    }
            
            # 构建任务信息
            task_info = {
                "任务标题": task.dimension_name,
                "任务要求": task_requirements
            }
            
            # 调用AI模型生成SQL（现在包含提示词生成逻辑）
            try:
                sql_response = await run_in_threadpool(
                    self.ai_advisor.generate_sql_from_task_prompt,
                    task_info,
                    dataset_schema
                )
                
                # 解析AI响应，提取SQL语句
                sql_statement = sql_response.get("sql", "")
                sql_explanation = sql_response.get("sql中文解释", "")
                status = sql_response.get("status", "failed")
                
                if status == "success" and sql_statement:
                    # 执行SQL获取数据
                    try:
                        # 使用DatasetService执行SQL获取数据
                        sql_result = await DatasetService.execute_sql_from_source(
                            dataset.data_source_id, 
                            sql_statement, 
                            session
                        )
                        
                        # 检查SQL执行是否成功
                        if sql_result.get("status"):
                            # SQL执行成功，获取数据并根据数据量处理
                            data = sql_result.get("data", [])
                            data_count = len(data)
                            
                            # 检查是否需要AI分析解读
                            analysis_type = task_requirements.get("分析类型", "数据获取")
                            
                            # 根据数据量大小采用不同的处理策略
                            if data_count == 0:
                                # 数据量为0，直接存储空数据
                                task.json_data = DataHelp.to_json([])
                                task.analysis_result = "未获取到数据"
                            elif data_count <= 1000:
                                # 数据量小于等于5000行，直接存储完整数据
                                task.json_data = DataHelp.to_json(data)
                                task.analysis_result = f"成功获取 {data_count} 条数据，数据已完整保存。"
                            elif 1000 < data_count <= 50000:
                                # 数据量在5000-50000之间，使用EnhancedDataCleaner生成数据摘要
                                try:
                                    from app.datatk.data_clear import EnhancedDataCleaner
                                    import pandas as pd
                                    
                                    # 将数据转换为DataFrame
                                    df = pd.DataFrame(data)
                                    
                                    # 使用EnhancedDataCleaner生成数据摘要
                                    cleaner = EnhancedDataCleaner()
                                    data_summary = cleaner.generate_data_summary(df)
                                    
                                    # 存储数据摘要到json_data
                                    task.json_data = DataHelp.to_json(data_summary)
                                    
                                    # 生成说明信息
                                    task.analysis_result = f"""数据量较大（{data_count} 条记录），已生成数据摘要：
                                        - 数据维度：{df.shape[0]} 行 × {df.shape[1]} 列
                                        完整的数据摘要已保存。"""
                                    
                                except Exception as summary_error:
                                    # 数据摘要生成失败，记录错误但仍标记为成功
                                    task.json_data = DataHelp.to_json([])  # 空数据
                                    task.analysis_result = f"获取了 {data_count} 条数据，但生成数据摘要时出错：{str(summary_error)}"
                                    
                            else:
                                # 数据量大于50000行，不存储数据，只记录说明
                                task.json_data = DataHelp.to_json([])  # 空数据
                                task.analysis_result = f"""数据量过大（{data_count} 条记录），超出存储限制。
                                    - 查询已成功执行，确认数据源包含 {data_count} 条符合条件的记录
                                    - 由于数据量较大，未保存具体数据内容
                                    - 建议：可通过增加筛选条件来减少数据量，或直接在数据集中查看结果"""
                            
                            # 更新任务状态为完成
                            task.analysis_sql = sql_statement
                            await self._update_task_status_safely(task, self.STATUS_COMPLETED, session)
                                                        
                            return {
                                "success": True,
                                "task_id": task.id,
                                "data_count": data_count,
                                "message": f"成功为任务 {task.dimension_name} 生成SQL并处理数据（{data_count} 条记录）" + 
                                          ("，并完成AI分析" if analysis_type == "数据分析" else "")
                            }
                        else:
                            # SQL执行失败，任务标记为失败
                            error_msg = sql_result.get("message", "SQL执行失败")
                            task.analysis_sql = sql_statement
                            task.analysis_result = sql_explanation
                            await self._update_task_status_safely(task, self.STATUS_FAILED, session)
                            task.analysis_result = sql_explanation                            
                            return {
                                "success": False,
                                "task_id": task.id,
                                "error": f"SQL执行失败: {error_msg}"
                            }
                            
                    except Exception as sql_error:
                        # SQL执行出错，任务标记为失败
                        task.analysis_sql = sql_statement
                        task.analysis_result = sql_explanation
                        await self._update_task_status_safely(task, self.STATUS_FAILED, session)
                        
                        return {
                            "success": False,
                            "task_id": task.id,
                            "error": f"执行SQL时出错: {str(sql_error)}"
                        }
                else:
                    # AI生成SQL失败，任务标记为失败
                    error_msg = sql_response.get("说明", "AI转换SQL失败")
                    task.analysis_result = error_msg
                    await self._update_task_status_safely(task, self.STATUS_FAILED, session)
                    return {
                        "success": False,
                        "task_id": task.id,
                        "error": f"AI生成SQL失败: {error_msg}"
                    }
            
            except Exception as e:
                # AI调用异常，任务标记为失败
                await self._update_task_status_safely(task, self.STATUS_FAILED, session)
                
                return {
                    "success": False,
                    "task_id": task.id,
                    "error": f"AI生成SQL异常: {str(e)}"
                }
            
        except Exception as e:
            # 更新任务状态为失败
            await self._update_task_status_safely(task, self.STATUS_FAILED, session)
            
            return {
                "success": False,
                "task_id": task.id,
                "error": f"执行数据查询任务异常: {str(e)}"
            }
    
    async def _update_task_status_safely(
        self, 
        task: AnalysisDimension, 
        status: str, 
        session: AsyncSession
    ) -> None:
        """
        安全地更新任务状态，包含完善的异常处理
        
        参数:
            task (AnalysisDimension): 任务对象
            status (str): 新状态
            session (AsyncSession): 数据库会话
        """
        try:
            task.status = status
            task.updated_at = datetime.now()
            session.add(task)
            await session.commit()
        except Exception as e:
            try:
                if session.in_transaction():
                    await session.rollback()
            except Exception as rollback_error:
                pass 